import os
import cv2
import shutil

# 视频抽帧
def extract_frames():
    # 定义文件夹路径
    video_dir = "Video"
    image_dir = "datasets/images"

    # 检查视频文件夹是否存在
    if not os.path.exists(video_dir):
        print(f"错误：{video_dir} 文件夹不存在")
        return

    # 清空并创建图片文件夹
    if os.path.exists(image_dir):
        shutil.rmtree(image_dir)
    os.makedirs(image_dir, exist_ok=True)

    # 获取所有MP4文件并按名称排序
    video_files = sorted([f for f in os.listdir(video_dir) if f.endswith(".mp4")])

    frame_counter = 1  # 全局帧计数器

    for video_file in video_files:
        video_path = os.path.join(video_dir, video_file)
        cap = cv2.VideoCapture(video_path)

        if not cap.isOpened():
            print(f"警告：无法打开文件 {video_file}，已跳过")
            continue

        print(f"正在处理: {video_file}...")

        # 获取视频帧率
        fps = cap.get(cv2.CAP_PROP_FPS)
        if fps <= 0:
            print(f"警告：视频 {video_file} 的帧率无效，将使用默认值 30 FPS")
            fps = 30.0

        # 计算每0.5秒对应的帧间隔
        frame_interval = max(1, int(round(fps * 0.5)))
        frame_index = 0

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            # 每0.5秒抽取一帧 (或达到间隔时)
            if frame_index % frame_interval == 0:
                # 生成带序号的文件名
                filename = f"P{frame_counter:05d}.jpg"
                output_path = os.path.join(image_dir, filename)

                # 保存帧图像
                if not cv2.imwrite(output_path, frame):
                    print(f"警告：未能保存 {filename}")
                else:
                    frame_counter += 1  # 仅当保存成功时增加计数器

            frame_index += 1

        cap.release()

    print(f"处理完成！共保存 {frame_counter - 1} 张图片到 {image_dir} 文件夹")


if __name__ == "__main__":
    extract_frames()